# 文件路径: src/data_access/timescale_repository.py
import random
import time
from datetime import datetime
from typing import List
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

from .base_repository import BaseRepository
from ..enums.timeframe import Timeframe


# 导入我们的抽象基类


class TimescaleRepository(BaseRepository):
    """
    使用 SQLAlchemy 和 Psycopg2 与 TimescaleDB 交互的具体仓储实现。
    """

    def __init__(self, db_url: str):
        """
        初始化仓储并创建数据库连接引擎。

        Args:
            db_url: SQLAlchemy 数据库连接字符串。
                    例如: "postgresql://user:password@host:port/database"
        """
        try:
            self.engine: Engine = create_engine(db_url)
            # 测试连接
            with self.engine.connect():
                print("成功连接到 TimescaleDB。")
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise

    def get_klines(
            self,
            symbol: str,
            timeframe: Timeframe,
            start_date: datetime,
            end_date: datetime
    ) -> pd.DataFrame:
        """
        从 TimescaleDB 中获取 K 线数据的具体实现。
        """
        # 使用参数化查询来防止 SQL 注入
        # 注意："time" 是 SQL 的关键字，最好用双引号括起来
        query = text("""
            SELECT *
            FROM klines
            WHERE
                symbol = :symbol AND
                timeframe = :timeframe AND
                "time" >= :start_date AND
                "time" < :end_date
            ORDER BY
                "time";
        """)

        try:
            df = pd.read_sql_query(
                sql=query,
                con=self.engine,
                params={
                    "symbol": symbol,
                    "timeframe": timeframe.value,
                    "start_date": start_date,
                    "end_date": end_date
                },
                index_col='time',
                parse_dates=['time']
            )
            return df
        except Exception as e:
            print(f"从数据库读取 K 线数据时出错: {e}")
            # 在实际应用中，你可能希望记录日志并返回一个空的 DataFrame
            return pd.DataFrame()

    def get_all_symbols(self) -> List[str]:
        """
        获取所有唯一交易对的具体实现。
        """
        query = text("SELECT DISTINCT symbol FROM klines;")
        try:
            with self.engine.connect() as connection:
                result = connection.execute(query)
                symbols = [row[0] for row in result]
                return symbols
        except Exception as e:
            print(f"获取所有 symbols 时出错: {e}")
            return []

    def save_klines(self, klines_df: pd.DataFrame):
        """
        高效且安全地保存 K 线数据，自动处理重复项。
        使用 INSERT ... ON CONFLICT 策略来防止重复插入。
        """
        if klines_df.empty:
            print("提供的 DataFrame 为空，无需保存。")
            return

        # 为临时表生成一个唯一的名称，以防多线程冲突
        temp_table_name = f"temp_klines_{int(time.time())}_{random.randint(0, 10000)}"

        # 将索引 'time' 重置为一列，以便写入
        df_to_write = klines_df.reset_index()

        # 确保列的顺序是确定的
        # (这只是一个好习惯，确保列顺序不会意外改变)
        cols_in_order = ['time', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume', 'quote_volume',
                         'trade_count', 'taker_buy_volume', 'taker_buy_quote_volume', 'close_time_raw', 'source']
        # 筛选出DataFrame中存在的列
        final_cols = [col for col in cols_in_order if col in df_to_write.columns]
        df_to_write = df_to_write[final_cols]

        # 使用事务来确保操作的原子性
        # with engine.begin() as conn: 实现了如果中途出错，所有操作都会回滚
        with self.engine.begin() as conn:
            try:
                # 1. 将数据写入临时表
                df_to_write.to_sql(
                    name=temp_table_name,
                    con=conn,
                    if_exists='replace',  # 临时表总是替换
                    index=False
                )

                # 2. 构造智能 INSERT 语句
                #    使用双引号来处理可能包含大写字母或关键字的列名
                cols_str = ", ".join(f'"{col}"' for col in final_cols)

                insert_sql = f"""
                    INSERT INTO klines ({cols_str})
                    SELECT {cols_str} FROM {temp_table_name}
                    ON CONFLICT (time, symbol, timeframe) DO NOTHING;
                """

                # 3. 执行 SQL
                conn.execute(text(insert_sql))

                print(f"成功处理 {len(df_to_write)} 条 K 线数据（重复项已自动跳过）。")

            except Exception as e:
                print(f"保存 K 线数据时出错: {e}")
                # 事务会自动回滚
                raise
            # 临时表会在事务结束后自动处理，但显式删除更干净
            # 但由于 `to_sql` 的 `if_exists='replace'`，其实风险很小

    def get_kline_timestamps(
            self,
            symbol: str,
            timeframe: str,
            start_date: datetime,
            end_date: datetime
    ) -> pd.DatetimeIndex:
        """
        从 TimescaleDB 中高效地只获取 K 线时间戳的具体实现。
        """
        query = text("""
            SELECT "time"
            FROM klines
            WHERE
                symbol = :symbol AND
                timeframe = :timeframe AND
                "time" >= :start_date AND
                "time" < :end_date
            ORDER BY "time";
        """)
        try:
            # read_sql_query 只查询一列并设为索引时，会返回一个只有索引的DataFrame
            df = pd.read_sql_query(
                sql=query,
                con=self.engine,
                params={
                    "symbol": symbol,
                    "timeframe": timeframe,
                    "start_date": start_date,
                    "end_date": end_date
                },
                index_col='time',
                parse_dates=['time']
            )
            return df.index
        except Exception as e:
            print(f"从数据库读取时间戳时出错: {e}")
            return pd.DatetimeIndex([])  # 出错时返回一个空的索引
